#!/usr/bin/env python
###############################################################################
# task.py: General TASK class
###############################################################################
# Copyright (C) 2009-10 CERN
#
# Authors: Jarno Rantala <jarno.rantala AT gmail.com>
#          Rohit Yadav <rohityadav89 AT gmail.com>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston MA 02110-1301, USA.
###############################################################################

import sys, string, xmlrpclib, time, subprocess, signal, os, traceback
from xml.dom import minidom
from unzip import *
from logger import *
from boinc import *

class TaskType:
    unknown, unzip_task, task, vm_task = range(4)

class TASK:
   taskType = TaskType.task # The type of task, can be task, unzip_task, vm_task
   appType  = "host_app" # Three types: host_app, boincvm_app, vm_app
   vmName = ""
   vmImagePath = ""
   application = ""
   copy_app_to_VM = 0
   copy_files_to_VM = []
   copy_stdin_to_VM = 0
   stdin_filename = ""
   stdout_filename = ""
   stderr_filename = ""
   copy_files_from_VM = []  
   checkpoint_filename = ""
   command_line = ""
   weight = 1.0
   CmdId = ""
   CmdResults = None
   app_process = None  # instance of Popen class of subprocess Module  
   suspended = 0
   starting_cpu = 0.0
   final_cpu_time = 0.0
   time_checkpointed = 0            
   # contribution of this task to overall fraction done
   final_cpu_time = 0;
   starting_cpu = 0;
   ready = 0;
   exitCode = None

   # how much CPU time was used by tasks before this in the job file
   suspended = 0 # zero or nonzero (false or true)
   wall_cpu_time = 0


   def readTag(self, tag, data):
      if tag == "virtualmachine":
         self.vmName = data
   
      elif tag == "image_path":
         self.vmImagePath = data
     
      elif tag == "copy_file_to_VM":
         self.copy_files_to_VM.append(data)
     
      elif tag == "application":
         self.application = data

      elif tag == "type":
         self.appType = data
     
      elif tag == "copy_app_to_VM":
         self.copy_app_to_VM = int(data)

      elif tag == "copy_stdin_to_VM":
         self.copy_stdin_to_VM = int(data)
 
      elif tag == "stdin_filename":
         self.stdin_filename = data

      elif tag == "stdout_filename":
         self.stdout_filename = data

      elif tag == "stderr_filename":
         self.stderr_filename = data

      elif tag == "copy_file_from_VM":
         self.copy_files_from_VM.append(data)
     
      elif tag == "command_line":
         self.command_line = data

      elif tag == "checkpoint_filename":
         self.checkpoint_filename = int(data)

      elif tag == "weight":
         self.weight = int(data)
     
      else:
         logger("Unknown Tag", tag)
   
   # Replace file names in command line with the physical names
   # resolved by boinc_resolve_filename-method. Every word which
   # starts "./" is recognised as file name.
   def resolve_commandline(self):
      newline = ""
      for word in self.command_line.split():
         if word[0:2] == "./":
            newline = newline + " " + boinc_resolve_filename(word)
         else:
            newline = newline + " " + word
      self.command_line = newline      

   def kill(self, vmController = None):
      if self.taskType == TaskType.vm_task:
         logger("VM", "Saving State")
         vmController.saveState(self.vmName) # saves and power off the VM
      else:
         if not self.ready:
            logger("Killing Process", self.command_line )
            self.app_process.kill()

   def stop(self, VMmanage):
      self.suspended = 1
      if self.vmName != "":
         VMmanage.pause(self.vmName)  
      else:
         self.app_process.send_signal(signal.SIGSTOP)

   def resume(self, VMmanage):
      self.suspended = 0
      if self.vmName != "":
         VMmanage.unpause(self.vmName)
      else:
         self.app_process.send_signal(signal.SIGCONT)

   def has_checkpointed(self, VMmanage, checkpoint_period):
     
      if self.vmName != "":
         if time.time()-self.time_checkpointed > checkpoint_period:
            # time to checkpoint
            VMmanage.saveSnapshot(self.vmName, self.checkpoint_filename)
            sys.stderr.write("snapshot at time: "+str(time.time())+"\n")
            self.time_checkpointed = time.time()  
            return 1
         else:
            return 0
      else:
         changed = 0
         if self.checkpoint_filename == "":
            return 0
         # is the file changed  ??

   def cpu_time(self, VMmanage):
      if self.vmName != "":
         # linux VM assumed!!!!
         if self.suspended: # we cannot send a cmd to VM which is paused
            return 0

         cmdid = VMmanage.runCmd(self.vmName, "cat", ["/proc/uptime"])

         wait = 1
         tic = time.time()
         while wait:
            for cmd in VMmanage.listFinishedCmds():
               if cmd == cmdid:
                  wait = 0
                  break
            # sys.stderr.write("Wait cpu_time cmd: "+str(wait)+"\n")
            time.sleep(1)
            if time.time() - tic > 10:
               sys.stderr.write("It took too long to get cpu time! \n")
               wait = 0

         res = VMmanage.getCmdResults(cmdid)['out']        
         return reduce( lambda x,y: x-y,  map(float, res.split()) )

      else:
         # cpu of subprocess  
         # sys.stderr.write(str(os.times()[2]) +"\n")    
         return os.times()[2]

   def poll(self, VMmanage = ""):
     
      if self.vmName != "":
         for Cmd in VMmanage.listFinishedCmds():
            if Cmd == self.CmdId:
               self.ready = 1
         if self.ready:
            self.CmdResults = VMmanage.getCmdResults(self.CmdId)
            self.exitCode = self.CmdResults['exitCodeOrSignal']
            self.final_cpu_time = self.CmdResults['resources']['ru_stime'] + self.CmdResults['resources']['ru_utime']

      else:
         self.exitCode = self.app_process.poll()
         if self.exitCode != None:
            self.ready = 1
            self.final_cpu_time = self.cpu_time(VMmanage)

   def VMrunning(self, VMmanage):
      running = 0
      for VM in VMmanage.listRunningVMs():
         if VM == self.vmName:
            running = 1
      return running
   
   def runVM(self, VMmanage, commandline = "", max_wait_time = float('inf')):
      app_path = boinc_resolve_filename(self.application)
      image_path = boinc_resolve_filename(self.vmImagePath)
      vm_path = boinc_resolve_filename(self.vmName)
      input_path = boinc_resolve_filename(self.stdin_filename)

      # Append wrapper's command-line arguments to those in the job file.
      self.command_line = self.command_line + " " + commandline#+ " < "+ self.stdin_filename

      # Check if the virtual machine is already on client
      doCreateVM = 1
     
      for VM in VMmanage.listAvailableVMs():
         if VM == self.vmName:      
            doCreateVM = 0
     
      if doCreateVM:
         # we assume that base directory of createVM is home_of_boinc/.VirtualBox
         # and that VM is in project directory home_of_boinc/projects/URL_of_project
         # in image_path we have the path of image relative to boinc/slot/n/ directory
         # that's why we have to remove first "../" and then it should be OK.
         image_path = image_path[3:len(image_path)]
         try:
            VMmanage.createVM(self.vmName, image_path)
         except Exception as e:
            sys.stderr.write("Creation of VM failed! \n")
            sys.stderr.write(str(e) + "\n")
            raise Exception(3)          

      # restore snapshot if there is one
      if VMmanage.getState(self.vmName) != "Saved":
         try:
            VMmanage.restoreSnapshot(self.vmName)
         except:
            sys.stderr.write("Restoring snapshot failed. \n")          

      # start VM
      self.time_checkpointed = time.time()
      if not self.VMrunning(VMmanage):
         try:
            VMmanage.start(self.vmName)
         except Exception as e:
            sys.stderr.write("Can't start VM: " + self.vmName + "\n")
            sys.stderr.write(str(e) + "\n")
            raise Exception(3)

      # wait until VM is running
      tic = time.time()
      while not self.VMrunning(VMmanage):
         sys.stderr.write("Wait VM '" + self.vmName + "' to run.\n")
         
         if time.time() - tic > max_wait_time:
            sys.stderr.write("It took too long to VM to run. \n")
            raise Exception(5)  

         time.sleep(1)

      # wait until VM has connected to broker successfully
      tic = time.time()
      while 1:
         try:
            VMmanage.ping(self.vmName)
         except Exception as e:
            sys.stderr.write("waiting for VM to connect to the broker \n")
               
            if time.time() - tic > max_wait_time:
               sys.stderr.write(str(e)+"\n")
               sys.stderr.write("It took too long to VM to connect to the broker. \n")
               raise Exception(5)

            time.sleep(10)
            continue
         break    

      # copy app file to VM
      if self.copy_app_to_VM:
         sys.stderr.write("copy file '"+app_path+"' to VM\n")
         
         [out, err, status] = VMmanage.cpFileToVM(self.vmName, app_path, self.application)
         if status:
            sys.stderr.write("Can't copy app to VM \n")
            sys.stderr.write(err)
            raise Exception(4)  
     
      # copy files file to VM
      for fileName in self.copy_files_to_VM:
         file_path = boinc_resolve_filename(fileName)
         sys.stderr.write("copy file '"+file_path+"' to VM\n")
         
         [out, err, status] = VMmanage.cpFileToVM(self.vmName, file_path, fileName)
         if status:
            sys.stderr.write("Can't copy files to VM \n")
            sys.stderr.write(err)
            raise Exception(4)  


      if self.application != "" or self.command_line != "":  # we put " " to command_line couple of lines above    
         sys.stderr.write("run command on VM: "+self.application+" "+self.command_line+"\n")      
         tic = time.time()
         while 1:
            try:
               self.CmdId = VMmanage.runCmd(self.vmName, self.application, [self.command_line], '{}',                                                                   'None', input_path)
            except Exception as e:
               sys.stderr.write("command didn't succeed.. try again.. \n")
 
               if time.time() - tic > max_wait_time:
                  sys.stderr.write(str(e)+"\n")
                  sys.stderr.write("Running command in VM didn't succeed. \n")
                  raise Exception(5)

               time.sleep(10)
               continue
            break    
       

   def run(self, commandline = ""):
      # unzip_task handled here
      if(self.taskType == TaskType.unzip_task):
         zipTool = unzip(True)
         try:
            zipTool.extract(self.command_line, self.stdout_filename)
         except IOError:
            logger("Error extracting: ", self.command_line)
         return

      app_path = boinc_resolve_filename(self.application)

      stdout_file = None
      stdin_file  = None
      stderr_file = None

      if self.stdout_filename != "":
         output_path = boinc_resolve_filename(self.stdout_filename)
         logger("stdout file", output_path)
         stdout_file = open(output_path, "a")

      if self.stdin_filename != "":      
         input_path = boinc_resolve_filename(self.stdin_filename)
         logger("stdin file", input_path)
         stdin_file = open(input_path, "r")
 
      if self.stderr_filename != "":
         err_path = boinc_resolve_filename(self.stderr_filename)
         logger("stderr file", err_path)
         stderr_file = open(err_path, "a")      
         
      # Append wrapper's command-line arguments to those in the job file.
      self.command_line = self.command_line + " " + commandline

      # resolve file names in command line
      self.resolve_commandline()
 
      logger("Wrapper Running ", app_path + " (" + self.command_line + ")")
           
      # runs application on host machine    
      self.app_process = subprocess.Popen((app_path+" "+self.command_line).split(), 0, None,        
                                                     stdin_file, stdout_file, stderr_file)  

   def debug(self):
      print '[TASK]', self.taskType, self.appType, self.vmName, self.application, self.command_line

